import torch
import os
import matplotlib.pyplot as plt
from tqdm import tqdm

from model.mesh_dataset import MeshDataset
from model.model_point_picker import ModelPointPicker

# 配置环境变量
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 定义项目根路径
PROJECT_ROOT = r"C:\Users\admin\PycharmProjects\Mesh-simplification-GNN"

# 定义训练参数
number_neigh_tri = 20
torch_dataset = MeshDataset(os.path.join(PROJECT_ROOT, "3d_models/stl/"))  # 使用全局路径拼接
device = torch.device("xpu" if torch.xpu.is_available() else ("cuda" if torch.cuda.is_available() else "cpu"))

# 加载模型
model_path = os.path.join(PROJECT_ROOT, 'save_models', 'point_picker', '1', '180.pt')  # 拼接模型路径
gnn_model = ModelPointPicker().to(device)
gnn_model.load_state_dict(torch.load(model_path))
gnn_model.eval()

# 获取数据集中的第一个图形数据
points = torch_dataset[0][0].x.cpu()

# 可视化生成的三维点
for i in range(0, 300, 10):
    fig = plt.figure()

    # 绘制3D散点图
    s = [1 for i in range(len(points))]
    ax = fig.add_subplot(121, projection='3d')
    ax.view_init(elev=110 + i, azim=-90)
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=s)

    ax = fig.add_subplot(122, projection='3d')
    ax.view_init(elev=110 + i, azim=-90)
    ax.scatter(points[:, 0], points[:, 1], points[:, 2])

    # 调整布局
    plt.tight_layout()
    plt.show(block=False)
    plt.pause(2)
    plt.close()

# 如果要可视化多个模型，解注释下面的代码
# saved_model_filenames = [f for f in os.listdir(os.path.join(PROJECT_ROOT, 'save_models', 'point_picker', '1')) if f.endswith('.pt')]
# for save_model in tqdm(saved_model_filenames):
#     model_path = os.path.join(PROJECT_ROOT, 'save_models', 'point_picker', '1', save_model)  # 拼接模型路径
#     gnn_model = ModelPointPicker().to(device)
#     gnn_model.load_state_dict(torch.load(model_path))
#     gnn_model.eval()

#     with torch.no_grad():
#         # 在此处进行模型推理和可视化操作
#         points = torch_dataset[0][0].x.cpu()
#         # 可视化部分代码同上
